source(here::here("R/00_setup.R"))

library(glue)
library(ggtext)

d_baseline <- read_csv(here("data/precinct_baseline.csv"), show_col_types = FALSE)
model <- read_rds(here("data/election_model.rds"))
d_control <- read_csv(here("data/control.csv"), show_col_types = FALSE)

d_st_base <- d_baseline |>
  group_by(state) |>
  summarize(baseline = log(sum(ndv)) - log(sum(nrv)))

# matrix of logit D vote shares
set.seed(2020)
natl_shift <- rnorm(100, 0, model$natl_sd)

# Apply precinct estimates to plan
d_seats <- map_dfr(
  c("OH-with-rules", "OH-no-rules",
    'AL-50-states', 'AL-Race-Blind', 'AL-Plaintiff-Frozen'), function(i) {

      abb <- str_sub(i, end = 2)
      # Select plans (either dataverse or simulated)
      if (i == "OH-with-rules") {
        plans <- load_50state_plans("OH")
      } else if (i == "OH-no-rules") {
        plans <- read_rds(here("data/ohio_remove-constraints.rds"))
      } else if (i == "AL-50-states") {
        plans <- load_50state_plans('AL')
      } else if (i == "AL-Race-Blind") {
        plans <- read_rds(here("data/alabama_raceblind.rds"))
      } else if (i == "AL-Plaintiff-Frozen") {
        plans <- read_rds(here("data/alabama_frozen_districts.rds"))
      }

      d_st <- filter(d_baseline, state == abb)
      ldvs <- part_dvs(plans, d_st, ndv, nrv) |>
        qlogis()
      seats <- outer(ldvs, natl_shift, `+`) |>
        model$prob_win_logit() |>
        rowMeans() |>
        matrix(ncol = 5001) |>
        colSums()

      tibble(
        state = abb,
        setting = switch(i,
                         `OH-with-rules` = "Our specification",
                         `OH-no-rules` = "Without state rules",
                         `AL-50-states` = "Our specification",
                         `AL-Race-Blind` = 'Without VRA race constraint',
                         `AL-Plaintiff-Frozen` = 'Plan with 2 VRA districts'
        ),
        seats_enac = seats[1],
        seats = seats[-1],
        draw = 1:5000
      )
    }) |>
  left_join(d_control, by = "state") |>
  left_join(d_st_base, by = "state")

label_seats <- function(x) {
  case_when(
    x == 0 ~ "No difference",
    x < 0 ~ number(-x, 1, prefix = "D+"),
    x > 0 ~ number(x, 1, prefix = "R+")
  )
}


# Main figure for endogeneity
d_seats_plot <- d_seats |>
  mutate(state = recode_factor(state, OH = "Ohio", AL = "Alabama")) |>
  mutate(setting = fct_relevel(setting, "Without state rules",
                               "Plan with 2 VRA districts",
                               "Our specification",
                               "Without VRA race constraint")) |>
  mutate(setting_preferred = setting == "Our specification",
         setting_lbl = fct_reorder(glue("<b style='color:{ifelse(setting_preferred, 'grey50', 'black')}'>{setting}</b>"),
                                   as.integer(setting)))

d_seats_plot |>
  ggplot(aes(x = seats - seats_enac, y = setting_lbl, color = factor(setting_preferred))) +
  geom_vline(xintercept = 0, lwd = 0.4) +
  stat_pointinterval(interval_size_range = c(0.8, 2.0), fatten_point = 1.0) +
  scale_x_continuous(
    "Difference of seats won under enacted plan\ncompared to baseline",
    expand = c(0, 0),
    labels = label_seats
  ) +
  ggforce::facet_col(state ~ ., scales = "free_y", space = "free") +
  coord_cartesian(xlim = c(-1, 2)) +
  labs(y = "Non-partisan baseline") +
  scale_color_manual(values = c(`FALSE` = "black", `TRUE` = "gray50")) +
  theme_bw() +
  guides(color = "none") +
  theme(axis.text.y=ggtext::element_markdown(),
        plot.margin = unit(c(0.1, 0.8, 0.1, 0.8), "cm"))

ggsave(here("paper/figures/endogeneity.pdf"), height = 3, width = 5)


# average by setting
d_seats |>
  group_by(state, setting) |>
  summarize(
    mean_diff = mean(seats - seats_enac),
    median_diff = median(seats - seats_enac)
  )

# within-draw differences across setting
d_seats_diff <- d_seats |>
  mutate(diff_seats = seats - seats_enac) |>
  pivot_wider(id_cols = c(draw, state), names_from = setting, values_from = diff_seats) |>
  janitor::clean_names() |>
  mutate(diff_draw = without_state_rules - our_specification)


quantile(d_seats_diff$diff_draw, c(0.025, 0.05, 0.10, 0.50, 0.90, 0.95, 0.975), na.rm = TRUE) |> round(3)
